import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv3D, Flatten, Dense, MaxPooling3D, GlobalAveragePooling3D
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import argparse


def run():
    # Load processed frame data and label data
    processed_frames = np.load('data.npy')
    processed_labels = np.load('labels.npy')

    # Split into training, validation, and test sets
    train_ratio = 0.6
    val_ratio = 0.2
    test_ratio = 0.2

    total_samples = processed_frames.shape[0]
    train_samples = int(total_samples * train_ratio)
    val_samples = int(total_samples * val_ratio)
    test_samples = total_samples - train_samples - val_samples

    train_frames = processed_frames[:train_samples]
    train_labels = processed_labels[:train_samples]
    val_frames = processed_frames[train_samples:train_samples + val_samples]
    val_labels = processed_labels[train_samples:train_samples + val_samples]
    test_frames = processed_frames[train_samples + val_samples:]
    test_labels = processed_labels[train_samples + val_samples:]

    print('Label Count (Label = 1):', np.count_nonzero(processed_labels == 1), 'Total Count:', len(processed_frames))
    print('Train Label Count (Label = 1):', np.count_nonzero(train_labels == 1), 'Total Count:', len(train_labels))
    print('Validation Label Count (Label = 1):', np.count_nonzero(val_labels == 1), 'Total Count:', len(val_labels))
    print('Test Label Count (Label = 1):', np.count_nonzero(test_labels == 1), 'Total Count:', len(test_labels))

    # Define the input shape of the model
    input_shape = processed_frames.shape[1:]

    # Create the model
    model = Sequential()
    model.add(Conv3D(64, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(1, 2, 2), padding='same'))
    model.add(Conv3D(64, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), padding='same'))
    model.add(Conv3D(128, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), padding='same'))
    model.add(Conv3D(128, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), padding='same'))
    model.add(Conv3D(256, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(2, 2, 2), padding='same'))
    model.add(Conv3D(256, 3, activation='relu', input_shape=input_shape, padding='same'))
    model.add(MaxPooling3D(pool_size=(2, 1, 1), padding='same'))
    model.add(GlobalAveragePooling3D())
    model.add(Dense(1, activation='sigmoid'))
    print(model.summary())
    # Compile the model
    model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

    # Define the callback to save the best model
    checkpoint = ModelCheckpoint(filepath='best_model.h5', monitor='val_loss', save_best_only=True)
    early_stop = EarlyStopping(monitor='val_loss', patience=args.early_stop_patience)

    # Train the model
    model.fit(train_frames, train_labels, validation_data=(val_frames, val_labels), batch_size=args.batch_size,
              epochs=args.epochs,
              callbacks=[checkpoint, early_stop])

    # Evaluate the model
    test_loss, test_acc = model.evaluate(test_frames, test_labels)
    print('Test Loss:', test_loss)
    print('Test Accuracy:', test_acc)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='This is an example of command-line arguments')

    # Add command-line arguments
    parser.add_argument('--folder_path', type=str,
                        default="forge_video_data")
    parser.add_argument('--epochs', type=int, default=500)
    parser.add_argument('--batch_size', type=int, default=128)
    parser.add_argument('--early_stop_patience', type=int, default=10)

    # Parse the command-line arguments
    args = parser.parse_args()
    run()
